#include "CudaParam.h"
#include <thrust/sequence.h>
#include <thrust/partition.h>
#include <ctime>

EpiModel* model = NULL;

GPUSimStruct gpu_data[NUM_GPUS];

vector<int>* host_comm_worker_start_vec = NULL;
vector<int>* host_comm_worker_end_vec = NULL;

int num_gpu_iter; // number of gpu iterations of calculating the day function
int num_gpu_ppl; // number of ppl on a gpu at one iteration
int num_outside_worker; // totoal number of outside workers

int getThreadsPerBlk(int threadsPerBlk){
	int largestPossible = 1024;
	while(largestPossible > 0){
		if(threadsPerBlk > largestPossible){
			threadsPerBlk = largestPossible;
			return threadsPerBlk;
		}
		largestPossible /= 2;
	}
	return 1;

}

void uploadVec(GPUSimStruct& data, int iter, int vec_id){
	
	int iter_first_comm = (*data.iter_comm_start)[iter];
	int iter_last_comm = (*data.iter_comm_end)[iter];

	int iter_first_ppl = model->commvec[iter_first_comm].nFirstPerson;
	int iter_last_ppl = model->commvec[iter_last_comm-1].nLastPerson;
	int iter_num_hommies = iter_last_ppl-iter_first_ppl;

	int iter_first_worker = (*host_comm_worker_start_vec)[iter_first_comm];
	int iter_last_worker = (*host_comm_worker_end_vec)[iter_last_comm-1];
	int iter_num_workers = iter_last_worker-iter_first_worker;

	// load person vec
	#pragma omp parallel for
	for(int i = iter_first_ppl; i < iter_last_ppl; i++){
		(*(data.host_ppl_vec[0]))[i - iter_first_ppl] = model->pvec[i];
	}

	// load worker vec
	#pragma omp parallel for
	for(int i = iter_first_comm; i < iter_last_comm; i++){
		int pos = (*host_comm_worker_start_vec)[i] - iter_first_worker;
		int size = model->commvec[i].workers.size();
		for(int j = 0; j < size; j++){
			(*(data.host_comm_worker_vec[0]))[pos+j] = model->pvec[model->commvec[i].workers[j]];
		}
	}

	// copy person vec to gpu
	cudaMemcpy(
		thrust::raw_pointer_cast(data.ppl_vec[vec_id]->data()),
		(&(*(data.host_ppl_vec[0]))[0]),
		iter_num_hommies*sizeof(Person),
		cudaMemcpyHostToDevice);

	// copy worker vec to gpu
	cudaMemcpy(
		thrust::raw_pointer_cast(data.comm_worker_vec[vec_id]->data()),
		(&(*(data.host_comm_worker_vec[0]))[0]),
		iter_num_workers*sizeof(Person),
		cudaMemcpyHostToDevice);
}

void downloadVec(GPUSimStruct& data, int iter, int vec_id){

	int iter_first_comm = (*data.iter_comm_start)[iter];
	int iter_last_comm = (*data.iter_comm_end)[iter];

	int iter_first_ppl = model->commvec[iter_first_comm].nFirstPerson;
	int iter_last_ppl = model->commvec[iter_last_comm-1].nLastPerson;
	int iter_num_hommies = iter_last_ppl-iter_first_ppl;

	int iter_first_worker = (*host_comm_worker_start_vec)[iter_first_comm];
	int iter_last_worker = (*host_comm_worker_end_vec)[iter_last_comm-1];
	int iter_num_workers = iter_last_worker-iter_first_worker;

	// copy person vec back to cpu
	cudaMemcpy(
		(&(*(data.host_ppl_vec[1]))[0]),
		thrust::raw_pointer_cast(data.ppl_vec[vec_id]->data()),
		iter_num_hommies*sizeof(Person),
		cudaMemcpyDeviceToHost);

	// copy worker vec back to gpu
	cudaMemcpy(
		(&(*(data.host_comm_worker_vec[1]))[0]),
		thrust::raw_pointer_cast(data.comm_worker_vec[vec_id]->data()),
		iter_num_workers*sizeof(Person),
		cudaMemcpyDeviceToHost);

	// probe person to epimodel
	#pragma omp parallel for
	for(int i = 0; i < iter_num_hommies; i++){
		Person& p = (*(data.host_ppl_vec[1]))[i];
		if(isInfected(p) && p.iday == -1)
			model->pvec[p.id] = p;
	}

	#pragma omp parallel for
	for(int i = 0 ; i < iter_num_workers; i++){
		Person& p = (*(data.host_comm_worker_vec[1]))[i];
		if(isInfected(p) && p.iday == -1)
			model->pvec[p.id] = p;
	}
}

void computeDay(GPUSimStruct& data, int iter, int vec_id){
	int iter_first_comm = (*data.iter_comm_start)[iter];
	int iter_last_comm = (*data.iter_comm_end)[iter];

	int iter_first_ppl = model->commvec[iter_first_comm].nFirstPerson;
	int iter_last_ppl = model->commvec[iter_last_comm-1].nLastPerson;
	int iter_num_hommies = iter_last_ppl-iter_first_ppl;

	int iter_first_worker = (*host_comm_worker_start_vec)[iter_first_comm];
	int iter_last_worker = (*host_comm_worker_end_vec)[iter_last_comm-1];
	int iter_num_workers = iter_last_worker-iter_first_worker;

	
	// fill the infection vector with all false
	thrust::fill(data.home_infect_result->begin(),data.home_infect_result->begin()+iter_num_hommies,false);
	thrust::fill(data.worker_infect_result->begin(),data.worker_infect_result->begin()+iter_num_workers,false);



	// setup the index and partition all people vec who is infectious
	thrust::sequence(data.home_ppl_index->begin(),data.home_ppl_index->begin()+iter_num_hommies);
	thrust::device_vector<int>::iterator homeInfectEndIterator = 
		thrust::partition(data.home_ppl_index->begin(),data.home_ppl_index->begin()+iter_num_hommies,
			IsHomeWorkerInfectious(thrust::raw_pointer_cast(data.ppl_vec[vec_id]->data())));

	// setup the index and partition all worker vec who is infectious
	thrust::sequence(data.worker_ppl_index->begin(),data.worker_ppl_index->begin()+iter_num_workers);
	thrust::device_vector<int>::iterator workInfectEndIterator = 
		thrust::partition(data.worker_ppl_index->begin(),data.worker_ppl_index->begin()+iter_num_workers,
			IsOutsideWorkerInfectious(thrust::raw_pointer_cast(data.comm_worker_vec[vec_id]->data())));


	int home_infect_num = homeInfectEndIterator - data.home_ppl_index->begin();
	if(home_infect_num != 0){
		int threadsPerBlk = num_gpu_ppl / home_infect_num;
		
		threadsPerBlk = getThreadsPerBlk(threadsPerBlk);

		DayOp3EvenBetter<<<home_infect_num,threadsPerBlk>>>(
			thrust::raw_pointer_cast(data.home_ppl_index->data()),
			thrust::raw_pointer_cast(data.ppl_vec[vec_id]->data()),
			thrust::raw_pointer_cast(data.tract_vec->data()),
			thrust::raw_pointer_cast(data.comm_vec->data()),
			thrust::raw_pointer_cast(data.rand_vec->data()),
			model->bTravel,
			thrust::raw_pointer_cast(data.comm_worker_vec[vec_id]->data()),
			thrust::raw_pointer_cast(data.comm_worker_start_vec->data()),
			thrust::raw_pointer_cast(data.comm_worker_end_vec->data()),
			iter_first_comm,
			thrust::raw_pointer_cast(data.home_infect_result->data()),
			thrust::raw_pointer_cast(data.worker_infect_result->data()),true);
	}


	int work_infect_num = workInfectEndIterator - data.worker_ppl_index->begin();
	if(work_infect_num != 0){
		int threadsPerBlk = num_gpu_ppl / work_infect_num;
		threadsPerBlk = getThreadsPerBlk(threadsPerBlk);

		DayOp3EvenBetter<<<work_infect_num,threadsPerBlk>>>(
			thrust::raw_pointer_cast(data.worker_ppl_index->data()),
			thrust::raw_pointer_cast(data.ppl_vec[vec_id]->data()),
			thrust::raw_pointer_cast(data.tract_vec->data()),
			thrust::raw_pointer_cast(data.comm_vec->data()),
			thrust::raw_pointer_cast(data.rand_vec->data()),
			model->bTravel,
			thrust::raw_pointer_cast(data.comm_worker_vec[vec_id]->data()),
			thrust::raw_pointer_cast(data.comm_worker_start_vec->data()),
			thrust::raw_pointer_cast(data.comm_worker_end_vec->data()),
			iter_first_comm,
			thrust::raw_pointer_cast(data.home_infect_result->data()),
			thrust::raw_pointer_cast(data.worker_infect_result->data()),false);
	}

	// partition all people who will be infected this period
	thrust::sequence(data.home_ppl_index->begin(),data.home_ppl_index->begin()+iter_num_hommies);
	homeInfectEndIterator =
		thrust::partition(data.home_ppl_index->begin(),data.home_ppl_index->begin()+iter_num_hommies,
		IsInfectedDay(thrust::raw_pointer_cast(data.home_infect_result->data())));

	// partition all worker who will be infected this period
	thrust::sequence(data.worker_ppl_index->begin(),data.worker_ppl_index->begin()+iter_num_workers);
	workInfectEndIterator = 
		thrust::partition(data.worker_ppl_index->begin(),data.worker_ppl_index->begin()+iter_num_workers,
		IsInfectedDay(thrust::raw_pointer_cast(data.worker_infect_result->data())));



	// real infection to hommies
	thrust::for_each(data.home_ppl_index->begin(),homeInfectEndIterator,
				InfectOp(thrust::raw_pointer_cast(data.rand_vec->data()),
				model->nTimer,model->AVEp,
				model->fSymptomaticAscertainment,model->bTrigger,
				model->fLiberalLeaveCompliance,model->fIsolationCompliance,
				thrust::raw_pointer_cast(data.comm_vec->data()),
				thrust::raw_pointer_cast(data.ppl_vec[vec_id]->data())));

	// real infection to workers
	thrust::for_each(data.worker_ppl_index->begin(),workInfectEndIterator,
				InfectOp(thrust::raw_pointer_cast(data.rand_vec->data()),
				model->nTimer,model->AVEp,
				model->fSymptomaticAscertainment,model->bTrigger,
				model->fLiberalLeaveCompliance,model->fIsolationCompliance,
				thrust::raw_pointer_cast(data.comm_vec->data()),
				thrust::raw_pointer_cast(data.comm_worker_vec[vec_id]->data())));
}